#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RTSP视频流接收器
该模块提供了一个用于接收和处理RTSP视频流的类
"""

import cv2
import numpy as np
import threading
import time
import logging
from typing import Optional, Tuple, Callable, Union, List, Dict, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('RTSPReceiver')


class RTSPReceiver:
    """
    RTSP视频流接收器类
    
    该类用于连接到RTSP视频流，读取视频帧，并提供各种控制和处理功能。
    
    属性:
        rtsp_url (str): RTSP流的URL
        buffer_size (int): 帧缓冲区大小
        reconnect_attempts (int): 连接断开时的重连尝试次数
        reconnect_delay (float): 重连尝试之间的延迟（秒）
    """
    
    def __init__(self, rtsp_url: str, buffer_size: int = 10, 
                 reconnect_attempts: int = 5, reconnect_delay: float = 2.0):
        """
        初始化RTSP接收器
        
        参数:
            rtsp_url (str): RTSP流的URL
            buffer_size (int, 可选): 帧缓冲区大小，默认为10
            reconnect_attempts (int, 可选): 连接断开时的重连尝试次数，默认为5
            reconnect_delay (float, 可选): 重连尝试之间的延迟（秒），默认为2.0
        """
        self.rtsp_url = rtsp_url
        self.buffer_size = buffer_size
        self.reconnect_attempts = reconnect_attempts
        self.reconnect_delay = reconnect_delay
        
        # 内部属性
        self._cap = None  # OpenCV VideoCapture对象
        self._is_running = False  # 指示接收器是否正在运行
        self._is_paused = False  # 指示接收器是否暂停
        self._frame_buffer = []  # 帧缓冲区
        self._current_frame = None  # 当前帧
        self._frame_count = 0  # 接收的帧计数
        self._last_frame_time = 0  # 上一帧的时间戳
        self._fps = 0  # 当前帧率
        self._lock = threading.Lock()  # 用于线程安全操作的锁
        self._thread = None  # 视频接收线程
        self._callbacks = []  # 帧处理回调函数列表
        self._connection_status = False  # 连接状态
        self._last_error = None  # 最后一个错误
        
    def connect(self) -> bool:
        """
        连接到RTSP流
        
        返回:
            bool: 连接成功返回True，否则返回False
        """
        try:
            logger.info(f"正在连接到RTSP流: {self.rtsp_url}")
            
            # 设置OpenCV的RTSP相关参数
            self._cap = cv2.VideoCapture(self.rtsp_url, cv2.CAP_FFMPEG)
            
            # 设置缓冲区大小
            self._cap.set(cv2.CAP_PROP_BUFFERSIZE, self.buffer_size)
            
            # 检查连接是否成功
            if not self._cap.isOpened():
                logger.error("无法连接到RTSP流")
                self._connection_status = False
                return False
            
            # 获取视频流信息
            self._width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            self._height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            self._fps = self._cap.get(cv2.CAP_PROP_FPS)
            
            logger.info(f"成功连接到RTSP流，分辨率: {self._width}x{self._height}, FPS: {self._fps}")
            self._connection_status = True
            return True
            
        except Exception as e:
            logger.error(f"连接RTSP流时发生错误: {str(e)}")
            self._last_error = str(e)
            self._connection_status = False
            return False
    
    def disconnect(self) -> None:
        """
        断开与RTSP流的连接
        """
        self.stop()
        if self._cap is not None:
            self._cap.release()
            self._cap = None
        self._connection_status = False
        logger.info("已断开与RTSP流的连接")
    
    def start(self) -> bool:
        """
        开始接收视频流
        
        返回:
            bool: 成功启动返回True，否则返回False
        """
        if self._is_running:
            logger.warning("接收器已经在运行")
            return True
        
        if not self._connection_status:
            success = self.connect()
            if not success:
                return False
        
        self._is_running = True
        self._is_paused = False
        self._thread = threading.Thread(target=self._receive_frames, daemon=True)
        self._thread.start()
        logger.info("开始接收视频流")
        return True
    
    def stop(self) -> None:
        """
        停止接收视频流
        """
        self._is_running = False
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=1.0)
        logger.info("停止接收视频流")
    
    def pause(self) -> None:
        """
        暂停接收视频流
        """
        self._is_paused = True
        logger.info("暂停接收视频流")
    
    def resume(self) -> None:
        """
        恢复接收视频流
        """
        self._is_paused = False
        logger.info("恢复接收视频流")
    
    def is_connected(self) -> bool:
        """
        检查是否已连接到RTSP流
        
        返回:
            bool: 已连接返回True，否则返回False
        """
        return self._connection_status
    
    def is_running(self) -> bool:
        """
        检查接收器是否正在运行
        
        返回:
            bool: 正在运行返回True，否则返回False
        """
        return self._is_running
    
    def is_paused(self) -> bool:
        """
        检查接收器是否已暂停
        
        返回:
            bool: 已暂停返回True，否则返回False
        """
        return self._is_paused
    
    def get_current_frame(self) -> Optional[np.ndarray]:
        """
        获取当前帧
        
        返回:
            Optional[np.ndarray]: 当前帧，如果没有可用帧则返回None
        """
        with self._lock:
            return self._current_frame.copy() if self._current_frame is not None else None
    
    def get_frame_info(self) -> Dict[str, Any]:
        """
        获取帧信息
        
        返回:
            Dict[str, Any]: 包含帧信息的字典
        """
        return {
            'width': self._width if hasattr(self, '_width') else None,
            'height': self._height if hasattr(self, '_height') else None,
            'fps': self._fps,
            'frame_count': self._frame_count,
            'is_running': self._is_running,
            'is_paused': self._is_paused,
            'connection_status': self._connection_status,
            'last_error': self._last_error
        }
    
    def add_frame_callback(self, callback: Callable[[np.ndarray], None]) -> None:
        """
        添加帧处理回调函数
        
        参数:
            callback (Callable[[np.ndarray], None]): 接收帧作为参数的回调函数
        """
        self._callbacks.append(callback)
        logger.info(f"添加了帧处理回调函数，当前回调函数数量: {len(self._callbacks)}")
    
    def remove_frame_callback(self, callback: Callable[[np.ndarray], None]) -> bool:
        """
        移除帧处理回调函数
        
        参数:
            callback (Callable[[np.ndarray], None]): 要移除的回调函数
            
        返回:
            bool: 成功移除返回True，否则返回False
        """
        if callback in self._callbacks:
            self._callbacks.remove(callback)
            logger.info(f"移除了帧处理回调函数，当前回调函数数量: {len(self._callbacks)}")
            return True
        return False
    
    def save_frame(self, filename: str, frame: Optional[np.ndarray] = None) -> bool:
        """
        保存帧为图像文件
        
        参数:
            filename (str): 文件名
            frame (Optional[np.ndarray], 可选): 要保存的帧，默认为当前帧
            
        返回:
            bool: 成功保存返回True，否则返回False
        """
        try:
            if frame is None:
                frame = self.get_current_frame()
                
            if frame is None:
                logger.error("没有可用的帧可保存")
                return False
                
            cv2.imwrite(filename, frame)
            logger.info(f"帧已保存到: {filename}")
            return True
            
        except Exception as e:
            logger.error(f"保存帧时发生错误: {str(e)}")
            self._last_error = str(e)
            return False
    
    def _receive_frames(self) -> None:
        """
        接收帧的内部方法（在单独的线程中运行）
        """
        reconnect_count = 0
        
        while self._is_running:
            try:
                # 如果暂停，则等待
                if self._is_paused:
                    time.sleep(0.1)
                    continue
                
                # 检查连接状态
                if not self._connection_status or self._cap is None:
                    if reconnect_count < self.reconnect_attempts:
                        logger.info(f"尝试重新连接 ({reconnect_count + 1}/{self.reconnect_attempts})")
                        success = self.connect()
                        if success:
                            reconnect_count = 0
                        else:
                            reconnect_count += 1
                            time.sleep(self.reconnect_delay)
                        continue
                    else:
                        logger.error(f"重连失败，已达到最大尝试次数: {self.reconnect_attempts}")
                        self._is_running = False
                        break
                
                # 读取帧
                ret, frame = self._cap.read()
                
                # 计算当前帧率
                current_time = time.time()
                if self._last_frame_time > 0:
                    time_diff = current_time - self._last_frame_time
                    if time_diff > 0:
                        self._fps = 0.8 * self._fps + 0.2 * (1.0 / time_diff)  # 平滑帧率
                self._last_frame_time = current_time
                
                if not ret:
                    logger.warning("无法读取帧，可能是流结束或连接问题")
                    self._connection_status = False
                    continue
                
                # 更新当前帧和帧计数
                with self._lock:
                    self._current_frame = frame
                    self._frame_count += 1
                    
                    # 更新帧缓冲区
                    if len(self._frame_buffer) >= self.buffer_size:
                        self._frame_buffer.pop(0)
                    self._frame_buffer.append(frame)
                
                # 处理回调函数
                for callback in self._callbacks:
                    try:
                        callback(frame.copy())
                    except Exception as e:
                        logger.error(f"执行帧回调函数时发生错误: {str(e)}")
                
            except Exception as e:
                logger.error(f"接收帧时发生错误: {str(e)}")
                self._last_error = str(e)
                self._connection_status = False
                time.sleep(0.1)  # 避免在错误情况下的快速循环
    
    def __enter__(self):
        """
        上下文管理器入口
        """
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口
        """
        self.disconnect()
    
    def __del__(self):
        """
        析构函数
        """
        self.disconnect()


# 示例用法
if __name__ == "__main__":
    # RTSP流URL示例
    rtsp_url = "rtsp://192.168.50.188:554"
    
    # 创建接收器实例
    receiver = RTSPReceiver(rtsp_url)
    
    try:
        # 连接并开始接收
        if receiver.connect():
            receiver.start()
            
            # 定义一个简单的帧处理回调函数
            def process_frame(frame):
                # 在这里可以添加自定义的帧处理逻辑
                # 例如：检测、识别、转换等
                pass
            
            # 添加回调函数
            receiver.add_frame_callback(process_frame)
            
            # 显示视频流
            window_name = "RTSP Stream"
            cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
            
            print("按 'q' 键退出")
            try:
                while True:
                    frame = receiver.get_current_frame()
                    if frame is not None:
                        cv2.imshow(window_name, frame)
                        
                    # 检查键盘输入
                    key = cv2.waitKey(1) & 0xFF
                    if key == ord('q'):
                        break
                    elif key == ord('s'):
                        # 按's'键保存当前帧
                        receiver.save_frame(f"frame_{receiver._frame_count}.jpg")
                    elif key == ord('p'):
                        # 按'p'键暂停/恢复
                        if receiver.is_paused():
                            receiver.resume()
                        else:
                            receiver.pause()
            finally:
                cv2.destroyAllWindows()
        else:
            print("无法连接到RTSP流")
    finally:
        # 确保资源被正确释放
        receiver.disconnect()